#include "node/image_visualization_node.h"

namespace project_template {
ImageVisualizationNode::ImageVisualizationNode() { init(); }

ImageVisualizationNode::~ImageVisualizationNode() {
  if (ros::isStarted()) {
    ros::shutdown();
    ros::waitForShutdown();
  }
  wait();
}

bool ImageVisualizationNode::init() {
  ros::NodeHandle nh;
  image_transport::ImageTransport it(nh);
  if (!ros::master::check()) {
    return false;
  }
  start();

  // 加载配置参数
  std::string source_dir = std::string(PTS_SOURCE_DIR);
  std::string image_param_path =
      source_dir + "/config/image_visualization_param.yaml";
  LoadConfigParam(image_param_path);

  sub_rgb_image_ =
      it.subscribe(image_visualization_config_.rgb_image_topic, 100,
                   &ImageVisualizationNode::RgbImageCallback, this);
  sub_depth_image_ =
      it.subscribe(image_visualization_config_.depth_image_topic, 100,
                   &ImageVisualizationNode::DepthImageCallback, this);

  return true;
}

void ImageVisualizationNode::run() {
  // 核心处理函数
  ros::Rate rate(30);

  while (ros::ok()) {
    ros::spinOnce();
    rate.sleep();
  }
  emit rosShutdown();
  quit();
}

void ImageVisualizationNode::RgbImageCallback(
    const sensor_msgs::ImageConstPtr& _msg) {
  cv_bridge::CvImagePtr cv_ptr;

  try {
    //深拷贝转换为opencv类型
    cv_bridge::CvImagePtr cv_ptr_compressed =
        cv_bridge::toCvCopy(_msg, sensor_msgs::image_encodings::BGR8);
    QImage img = Mat2QImage(cv_ptr_compressed->image);
    emit showRGBImage(img);
  } catch (cv_bridge::Exception& e) {
    return;
  }
}

void ImageVisualizationNode::DepthImageCallback(
    const sensor_msgs::ImageConstPtr& _msg) {
  cv_bridge::CvImagePtr cv_ptr;

  try {
    //深拷贝转换为opencv类型
    cv_bridge::CvImagePtr cv_ptr_compressed =
        cv_bridge::toCvCopy(_msg, sensor_msgs::image_encodings::TYPE_16UC1);
    QImage img = Mat2QImage(cv_ptr_compressed->image);
    emit showDepthImage(img);
  } catch (cv_bridge::Exception& e) {
    return;
  }
}

QImage ImageVisualizationNode::Mat2QImage(const cv::Mat& _src) {
  QImage dest(_src.cols, _src.rows, QImage::Format_ARGB32);

  const float scale = 255.0;

  if (_src.depth() == CV_8U) {
    if (_src.channels() == 1) {
      for (int i = 0; i < _src.rows; ++i) {
        for (int j = 0; j < _src.cols; ++j) {
          int level = _src.at<quint8>(i, j);
          dest.setPixel(j, i, qRgb(level, level, level));
        }
      }
    } else if (_src.channels() == 3) {
      for (int i = 0; i < _src.rows; ++i) {
        for (int j = 0; j < _src.cols; ++j) {
          cv::Vec3b bgr = _src.at<cv::Vec3b>(i, j);
          dest.setPixel(j, i, qRgb(bgr[2], bgr[1], bgr[0]));
        }
      }
    }
  } else if (_src.depth() == CV_32F) {
    if (_src.channels() == 1) {
      for (int i = 0; i < _src.rows; ++i) {
        for (int j = 0; j < _src.cols; ++j) {
          int level = scale * _src.at<float>(i, j);
          dest.setPixel(j, i, qRgb(level, level, level));
        }
      }
    } else if (_src.channels() == 3) {
      for (int i = 0; i < _src.rows; ++i) {
        for (int j = 0; j < _src.cols; ++j) {
          cv::Vec3f bgr = scale * _src.at<cv::Vec3f>(i, j);
          dest.setPixel(j, i, qRgb(bgr[2], bgr[1], bgr[0]));
        }
      }
    }
  } else if (_src.depth() == CV_16U) {
    if (_src.channels() == 1) {
      for (int i = 0; i < _src.rows; ++i) {
        for (int j = 0; j < _src.cols; ++j) {
          int level =
              scale *
              (_src.at<uint16_t>(i, j) /
               static_cast<float>(std::numeric_limits<uint16_t>::max()));
          dest.setPixel(j, i, qRgb(level, level, level));
        }
      }
    }
  }

  return dest;
}

cv::Mat ImageVisualizationNode::QImage2Mat(QImage& _image) {
  cv::Mat mat;
  switch (_image.format()) {
    case QImage::Format_ARGB32:
    case QImage::Format_RGB32:
    case QImage::Format_ARGB32_Premultiplied:
      mat = cv::Mat(_image.height(), _image.width(), CV_8UC4,
                    (void*)_image.constBits(), _image.bytesPerLine());
      break;
    case QImage::Format_RGB888:
      mat = cv::Mat(_image.height(), _image.width(), CV_8UC3,
                    (void*)_image.constBits(), _image.bytesPerLine());
      cv::cvtColor(mat, mat, CV_BGR2RGB);
      break;
    case QImage::Format_Indexed8:
      mat = cv::Mat(_image.height(), _image.width(), CV_8UC1,
                    (void*)_image.constBits(), _image.bytesPerLine());
      break;
  }
  return mat;
}

void ImageVisualizationNode::LoadConfigParam(std::string _path_of_yaml) {
  YAML::Node config;
  try {
    config = YAML::LoadFile(_path_of_yaml);
  } catch (...) {
    AERROR_F(
        "The format of config file [ %s ] is wrong, Please check (e.g. "
        "indentation).",
        _path_of_yaml.c_str());
    return;
  }

  for (size_t i = 0; i < config.size(); ++i) {
    image_visualization_config_ = LoadImageVisualizationConfig(config[i]);
  }
}

}  // namespace project_template
